"""
Configuration management for Easy RWKV
"""

import os
import yaml
from typing import Dict, Any, Optional
from dataclasses import dataclass, asdict


@dataclass
class Config:
    """Configuration class for Easy RWKV"""
    
    # Model settings
    model_path: str = ""
    strategy: str = "cuda fp16"
    state_path: Optional[str] = None
    
    # Generation settings
    temperature: float = 1.0
    top_p: float = 0.3
    alpha_presence: float = 0.5
    alpha_frequency: float = 0.5
    penalty_decay: float = 0.996
    max_tokens: int = 256
    chunk_len: int = 256
    
    # Environment settings
    rwkv_v7_on: bool = True
    rwkv_jit_on: bool = True
    rwkv_cuda_on: bool = False
    
    # Chat settings
    system_prompt: str = "User: hi\n\nAssistant: Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.\n\n"
    
    @classmethod
    def from_file(cls, config_path: str) -> "Config":
        """Load configuration from YAML file"""
        if not os.path.exists(config_path):
            raise FileNotFoundError(f"Config file not found: {config_path}")
        
        with open(config_path, 'r', encoding='utf-8') as f:
            config_dict = yaml.safe_load(f)
        
        return cls(**config_dict)
    
    def save_to_file(self, config_path: str) -> None:
        """Save configuration to YAML file"""
        os.makedirs(os.path.dirname(config_path), exist_ok=True)
        
        with open(config_path, 'w', encoding='utf-8') as f:
            yaml.dump(asdict(self), f, default_flow_style=False, allow_unicode=True)
    
    @classmethod
    def get_default_config_path(cls) -> str:
        """Get default configuration file path"""
        home_dir = os.path.expanduser("~")
        return os.path.join(home_dir, ".easy_rwkv", "config.yaml")
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert config to dictionary"""
        return asdict(self)
    
    def update(self, **kwargs) -> None:
        """Update configuration with new values"""
        for key, value in kwargs.items():
            if hasattr(self, key):
                setattr(self, key, value)
            else:
                raise ValueError(f"Unknown config key: {key}")
